#!/bin/env python

import MDAnalysis as mda
from MDAnalysis.analysis.distances import distance_array
import numpy as np
import matplotlib.pyplot as plt
from numpy.linalg import svd
import argparse
 
parser = argparse.ArgumentParser(description="分子轨迹.xyz")
parser.add_argument('-f',"--file",required=True,type=str, help="xyz文件")
parser.add_argument('-p',"--part",required=True,type=str, help="原子序号")
parser.add_argument('-l',"--line",required=True,type=str, help="线")

args = parser.parse_args()

# 你的轨迹文件和拓扑文件
xyz_file = args.file
group1_atoms_str,group2_atoms_str = args.part.split(';')
atom1_id,atom2_id,atom3_id,atom4_id = args.line.split(';')

atom1_idx = int(atom1_id)
atom2_idx = int(atom2_id)
atom3_idx = int(atom3_id)
atom4_idx = int(atom4_id)

print(f'part1:{group1_atoms_str}')
print(f'part2:{group2_atoms_str}')

# topology_file = 'your_topology.pdb'  # 使用你的拓扑文件

plt.rcParams['font.sans-serif'] = ['Arial']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def parse_atom_indices(atom_indices_str):
    indices = []
    for part in atom_indices_str.split(','):
        if '-' in part:
            start, end = part.split('-')
            indices.extend(range(int(start), int(end) + 1))
        else:
            indices.append(int(part))
    indices = np.array(indices) - 1   # MDAnalysis 与 Gaussian 中编号不同，需要修改
    return indices

# 解析原子编号
group1_atoms = parse_atom_indices(group1_atoms_str)
group2_atoms = parse_atom_indices(group2_atoms_str)

# 加载轨迹
u = mda.Universe(xyz_file)

# 选择两个部分
group1 = u.select_atoms(f'index {" ".join(map(str, group1_atoms))}')
group2 = u.select_atoms(f'index {" ".join(map(str, group2_atoms))}')

# 准备存储质心距离
distances = []

# 遍历轨迹中的每一帧
for ts in u.trajectory:
    # 计算两个部分的质心
    com1 = group1.center_of_mass()
    com2 = group2.center_of_mass()
    
    # 计算质心距离
    distance = np.linalg.norm(com1 - com2)
    distances.append(distance)

# 将距离转换为numpy数组
distances = np.array(distances)

# 计算平均值、标准差
mean_distance = np.mean(distances)
std_distance = np.std(distances)

print(f"Mean distance to the center of mass: {mean_distance:.3f} Å")
print(f"Standard deviation of distance from center of mass: {std_distance:.3f} Å")

# 绘制质心距离随时间变化的图
#plt.figure()
#plt.plot(distances)
#plt.xlabel('Time (frames)')
#plt.ylabel('Centroid distance(Å)')
#plt.title('The distance between the center of mass of the two parts varies with time')
#plt.show()

# 绘制质心距离的分布图
#plt.figure()
#plt.hist(distances, bins=30)
#plt.xlabel('Centroid distance (Å)')
#plt.ylabel('Frequency')
#plt.title('Centroid distance distribution')
#plt.show()

# 计算平面
def calculate_normal_vector(positions):
    """计算一组点的最小二乘平面的法向量"""
    centroid = np.mean(positions, axis=0)
    centered_positions = positions - centroid
    _, _, vh = svd(centered_positions)
    normal_vector = vh[-1]
    return normal_vector

def compute_dihedral_angle(n1, n2):
    """计算两个法向量之间的二面角"""
    cos_angle = np.dot(n1, n2) / (np.linalg.norm(n1) * np.linalg.norm(n2))
    angle = np.arccos(np.clip(cos_angle, -1.0, 1.0))  # 确保值在 [-1, 1] 范围内
    return np.degrees(angle)

# 准备存储二面角
dihedral_angles = []

# 遍历轨迹中的每一帧
for ts in u.trajectory:
    # 计算两个部分的最小二乘平面的法向量
    normal1 = calculate_normal_vector(group1.positions)
    normal2 = calculate_normal_vector(group2.positions)
    
    # 计算二面角
    angle = compute_dihedral_angle(normal1, normal2)
    if angle > 90 : 
	    angle = 180 - angle
    else :
          pass
    dihedral_angles.append(angle)

# 将角度转换为numpy数组
dihedral_angles = np.array(dihedral_angles)

# 计算平均值、标准差
mean_angle = np.mean(dihedral_angles)
std_angle = np.std(dihedral_angles)

# print(f"平均二面角: {mean_angle:.3f} 度")
# print(f"二面角标准差: {std_angle:.3f} 度")
print(f"Mean dihedral angle: {mean_angle:.3f} degrees")
print(f"Dihedral angle standard deviation: {std_angle:.3f} degrees")
# 绘制二面角随时间变化的图
#plt.figure()
#plt.plot(dihedral_angles)
# plt.xlabel('时间 (帧)')
# plt.ylabel('二面角 (度)')
# plt.title('两个部分的二面角随时间变化')
#plt.xlabel('Time (frames)')
#plt.ylabel('Dihedral Angle (degrees)')
#plt.title('Dihedral Angle Between Two Groups Over Time')
#plt.show()

# 绘制二面角的分布图
#plt.figure()
#plt.hist(dihedral_angles, bins=30)
# plt.xlabel('二面角 (度)')
# plt.ylabel('频数')
# plt.title('二面角分布')
#plt.xlabel('Dihedral Angle (degrees)')
#plt.ylabel('Frequency')
#plt.title('Distribution of Dihedral Angles')
#plt.show()





# 定义你感兴趣的原子组的编号
# 假设两个原子组分别为 group1 和 group2, 中点计算使用 atom1 和 atom2, 另一个中点计算使用 atom3 和 atom4

atom1 = u.atoms[atom1_idx - 1]  # 注意：索引从0开始
atom2 = u.atoms[atom2_idx - 1]
atom3 = u.atoms[atom3_idx - 1]
atom4 = u.atoms[atom4_idx - 1]

angles_line = []

# 迭代轨迹帧
for ts in u.trajectory:
    # 计算中点
    midpoint1 = 0.5 * (atom1.position + atom2.position)
    midpoint2 = 0.5 * (atom3.position + atom4.position)
    
    # 计算group1的质心 midpoint
    centroid1 = group1.centroid()
    
    # 计算平面法向量
    plane_vector1 = group1.centroid() - midpoint1
    plane_vector2 = group1.centroid() - midpoint2
	
    # normal_vector = np.cross(plane_vector1, plane_vector2)
    # 计算向量
    # line_vector = midpoint2 - midpoint1
    
    # 计算角度
    cos_theta = np.dot(plane_vector1, plane_vector2) / (np.linalg.norm(plane_vector1) * np.linalg.norm(plane_vector2))
    angle = np.degrees(np.arccos(np.clip(cos_theta, -1.0, 1.0)))
    
    # 将角度转换为 0-90 度范围
    if angle > 90:
        angle = 180 - angle
    
    angles_line.append(angle)

# 转换为 numpy 数组
angles_line = np.array(angles_line)
# print(angles_line)
# 计算统计数据
mean_angle_line = np.mean(angles_line)
std_angle_line = np.std(angles_line)

print(f'Mean Angle: {mean_angle_line:.2f} degrees')
print(f'Standard Deviation: {std_angle_line:.2f} degrees')

# 绘制角度随轨迹帧编号变化的图
#plt.figure()
#plt.plot(range(len(angles_line)), angles_line, label='Angle over Trajectory Frames')
#plt.xlabel('Trajectory Frame')
#plt.ylabel('Angle (degrees)')
#plt.title('Angle between Line and Plane over Trajectory Frames')
#plt.legend()
#plt.show()

# 绘制角度分布直方图
#plt.figure()
#plt.hist(angles_line, bins=30, alpha=0.75)
#plt.xlabel('Angle (degrees)')
#plt.ylabel('Frequency')
#plt.title('Distribution of Angles')
#plt.show()


output_filename = "distances_output.txt"
with open(output_filename, 'w') as f:
    f.write("Index \t distance (Å)\n")
    for idx, distances in enumerate(distances):
        f.write(f"{idx + 1}\t{distances:.2f}\n")

output_filename = "dihedral_angles_output.txt"
with open(output_filename, 'w') as f:
    f.write("Index \t Angle (degrees)\n")
    for idx, angle in enumerate(dihedral_angles):
        f.write(f"{idx + 1}\t{angle:.2f}\n")
		
output_filename = "angles_line_output.txt"
with open(output_filename, 'w') as f:
    f.write("Index \t Angle (degrees)\n")
    for idx, angle in enumerate(angles_line):
        f.write(f"{idx + 1}\t{angle:.2f}\n")



